#!/anaconda3/envs/FEALPy/bin python3.8
# -*- coding: utf-8 -*-
# File: get_bce_embedding.py
# Author: Bryan SHEN
# E-mail: m18801919240_3@163.com
# Site: Shanghai, China
# Time: 2024/4/22 11:59
# File-Desp:

import requests
import json
import pandas as pd
from tqdm import tqdm

# URL of the FastAPI endpoint
url = 'http://192.168.22.46:8091/bce_embeddings/'


# Function to get embeddings and save them with identifiers
def get_and_save_embeddings(sentences, embeddings_file):

    chunk_size = 64  # Adjust based on server configuration and request payload size
    with open(embeddings_file, 'w') as f:
        for i in tqdm(range(0, len(sentences), chunk_size)):
            response = requests.post(url, json={"sentences": sentences[i:i + chunk_size]})
            if response.status_code == 200:
                data = response.json()['embeddings']
                for idx, emb in enumerate(data, start=i):
                    json_record = json.dumps({"id": idx, "embedding": emb}, ensure_ascii=False)
                    f.write(json_record + '\n')
            else:
                print("Failed to get embeddings for chunk {}: {}".format(i // chunk_size, response.text))


if __name__ == '__main__':

    raw_data = "../../data/raw/聚类-0416-Endata.xlsx"
    df_raw = pd.read_excel(raw_data, sheet_name='4w热点')
    # df_raw["text"] = df_raw["热点"].str.cat(df_raw[['作品1', '作品2', '作品3', '作品4', '作品5']], sep='\n', na_rep='')  # 多列合并
    df_raw["text"] = df_raw["热点"]  # 只选取title
    sentences = df_raw["text"].tolist()
    # File to store the embeddings
    embeddings_file = '../../data/result/embeddings_title.jsonl'
    get_and_save_embeddings(sentences, embeddings_file)

